"""
线程管理模块
处理多线程和性能优化
"""

import asyncio
import concurrent.futures
import threading
import queue
import time
import logging
from typing import Dict, List, Any, Callable, Optional
from functools import wraps
import psutil
import gc

class ThreadManager:
    """线程管理器"""
    
    def __init__(self, config: Dict[str, Any]):
        """初始化线程管理器"""
        self.config = config
        self.logger = logging.getLogger('image_cropping_system')
        
        # 线程池配置
        threading_config = config.get('threading', {})
        self.max_workers = threading_config.get('max_workers', 8)
        self.io_pool_size = threading_config.get('io_pool_size', 4)
        self.compute_pool_size = threading_config.get('compute_pool_size', 4)
        
        # 创建线程池
        self.io_executor = concurrent.futures.ThreadPoolExecutor(
            max_workers=self.io_pool_size,
            thread_name_prefix="IO-Worker"
        )
        
        self.compute_executor = concurrent.futures.ThreadPoolExecutor(
            max_workers=self.compute_pool_size,
            thread_name_prefix="Compute-Worker"
        )
        
        # 任务队列
        self.task_queue = queue.Queue()
        self.result_queue = queue.Queue()
        
        # 性能监控
        self.performance_stats = {
            'total_tasks': 0,
            'completed_tasks': 0,
            'failed_tasks': 0,
            'average_processing_time': 0.0,
            'memory_usage': 0.0,
            'cpu_usage': 0.0
        }
        
        # 缓存配置
        cache_config = config.get('cache', {})
        self.cache_max_size = cache_config.get('max_size', 100)
        self.cache_ttl = cache_config.get('ttl', 3600)
        self.cache = {}
        self.cache_timestamps = {}
        
        self.logger.info(f"线程管理器初始化完成 - IO线程池: {self.io_pool_size}, 计算线程池: {self.compute_pool_size}")
    
    def submit_io_task(self, func: Callable, *args, **kwargs) -> concurrent.futures.Future:
        """提交I/O密集型任务"""
        future = self.io_executor.submit(self._wrap_task(func), *args, **kwargs)
        self.performance_stats['total_tasks'] += 1
        return future
    
    def submit_compute_task(self, func: Callable, *args, **kwargs) -> concurrent.futures.Future:
        """提交计算密集型任务"""
        future = self.compute_executor.submit(self._wrap_task(func), *args, **kwargs)
        self.performance_stats['total_tasks'] += 1
        return future
    
    def _wrap_task(self, func: Callable) -> Callable:
        """包装任务函数，添加性能监控"""
        @wraps(func)
        def wrapped_func(*args, **kwargs):
            start_time = time.time()
            
            try:
                # 检查内存使用
                self._check_memory_usage()
                
                # 执行任务
                result = func(*args, **kwargs)
                
                # 更新统计信息
                processing_time = time.time() - start_time
                self._update_performance_stats(processing_time, success=True)
                
                return result
                
            except Exception as e:
                # 更新错误统计
                processing_time = time.time() - start_time
                self._update_performance_stats(processing_time, success=False)
                
                self.logger.error(f"任务执行失败: {e}")
                raise
        
        return wrapped_func
    
    def _check_memory_usage(self) -> None:
        """检查内存使用情况"""
        try:
            process = psutil.Process()
            memory_info = process.memory_info()
            memory_mb = memory_info.rss / 1024 / 1024
            
            self.performance_stats['memory_usage'] = memory_mb
            
            # 如果内存使用过高，触发垃圾回收
            if memory_mb > 1024:  # 超过1GB
                gc.collect()
                self.logger.warning(f"内存使用过高 ({memory_mb:.2f}MB)，触发垃圾回收")
                
        except Exception as e:
            self.logger.error(f"检查内存使用时发生错误: {e}")
    
    def _update_performance_stats(self, processing_time: float, success: bool) -> None:
        """更新性能统计信息"""
        if success:
            self.performance_stats['completed_tasks'] += 1
        else:
            self.performance_stats['failed_tasks'] += 1
        
        # 更新平均处理时间
        total_completed = self.performance_stats['completed_tasks']
        if total_completed > 0:
            current_avg = self.performance_stats['average_processing_time']
            self.performance_stats['average_processing_time'] = (
                (current_avg * (total_completed - 1) + processing_time) / total_completed
            )
    
    async def batch_process(self, tasks: List[Dict[str, Any]], batch_size: int = 4) -> List[Any]:
        """批量处理任务"""
        results = []
        
        for i in range(0, len(tasks), batch_size):
            batch = tasks[i:i + batch_size]
            
            # 并行处理批次
            futures = []
            for task in batch:
                task_type = task.get('type', 'io')
                func = task['func']
                args = task.get('args', [])
                kwargs = task.get('kwargs', {})
                
                if task_type == 'compute':
                    future = self.submit_compute_task(func, *args, **kwargs)
                else:
                    future = self.submit_io_task(func, *args, **kwargs)
                
                futures.append(future)
            
            # 等待批次完成
            try:
                batch_results = await asyncio.gather(
                    *[asyncio.wrap_future(future) for future in futures],
                    return_exceptions=True
                )
                
                for result in batch_results:
                    if isinstance(result, Exception):
                        self.logger.error(f"批次任务失败: {result}")
                        results.append(None)
                    else:
                        results.append(result)
                        
            except Exception as e:
                self.logger.error(f"批次处理失败: {e}")
                results.extend([None] * len(batch))
            
            # 更新进度
            progress = min(i + batch_size, len(tasks))
            self.logger.info(f"批量处理进度: {progress}/{len(tasks)}")
        
        return results
    
    def get_from_cache(self, key: str) -> Optional[Any]:
        """从缓存获取数据"""
        try:
            if key in self.cache:
                timestamp = self.cache_timestamps.get(key, 0)
                if time.time() - timestamp < self.cache_ttl:
                    return self.cache[key]
                else:
                    # 缓存过期，删除
                    del self.cache[key]
                    del self.cache_timestamps[key]
        except Exception as e:
            self.logger.error(f"从缓存获取数据时发生错误: {e}")
        
        return None
    
    def set_cache(self, key: str, value: Any) -> None:
        """设置缓存数据"""
        try:
            # 检查缓存大小
            if len(self.cache) >= self.cache_max_size:
                self._cleanup_cache()
            
            self.cache[key] = value
            self.cache_timestamps[key] = time.time()
            
        except Exception as e:
            self.logger.error(f"设置缓存数据时发生错误: {e}")
    
    def _cleanup_cache(self) -> None:
        """清理过期缓存"""
        try:
            current_time = time.time()
            expired_keys = []
            
            for key, timestamp in self.cache_timestamps.items():
                if current_time - timestamp > self.cache_ttl:
                    expired_keys.append(key)
            
            for key in expired_keys:
                if key in self.cache:
                    del self.cache[key]
                if key in self.cache_timestamps:
                    del self.cache_timestamps[key]
            
            # 如果还是太大，删除最旧的缓存
            if len(self.cache) >= self.cache_max_size:
                sorted_items = sorted(self.cache_timestamps.items(), key=lambda x: x[1])
                items_to_remove = len(self.cache) - self.cache_max_size + 1
                
                for key, _ in sorted_items[:items_to_remove]:
                    if key in self.cache:
                        del self.cache[key]
                    if key in self.cache_timestamps:
                        del self.cache_timestamps[key]
                        
        except Exception as e:
            self.logger.error(f"清理缓存时发生错误: {e}")
    
    def get_performance_stats(self) -> Dict[str, Any]:
        """获取性能统计信息"""
        try:
            # 获取CPU使用率
            cpu_percent = psutil.cpu_percent(interval=1)
            self.performance_stats['cpu_usage'] = cpu_percent
            
            # 获取内存使用情况
            process = psutil.Process()
            memory_info = process.memory_info()
            self.performance_stats['memory_usage'] = memory_info.rss / 1024 / 1024
            
        except Exception as e:
            self.logger.error(f"获取性能统计信息时发生错误: {e}")
        
        return self.performance_stats.copy()
    
    def shutdown(self) -> None:
        """关闭线程池"""
        try:
            self.io_executor.shutdown(wait=True)
            self.compute_executor.shutdown(wait=True)
            self.logger.info("线程池已关闭")
        except Exception as e:
            self.logger.error(f"关闭线程池时发生错误: {e}")
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown()

class PerformanceMonitor:
    """性能监控器"""
    
    def __init__(self, thread_manager: ThreadManager):
        """初始化性能监控器"""
        self.thread_manager = thread_manager
        self.logger = logging.getLogger('image_cropping_system')
        self.monitoring = False
        self.monitor_thread = None
    
    def start_monitoring(self, interval: int = 30) -> None:
        """开始性能监控"""
        if self.monitoring:
            return
        
        self.monitoring = True
        self.monitor_thread = threading.Thread(
            target=self._monitor_loop,
            args=(interval,),
            daemon=True
        )
        self.monitor_thread.start()
        self.logger.info("性能监控已启动")
    
    def stop_monitoring(self) -> None:
        """停止性能监控"""
        self.monitoring = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=5)
        self.logger.info("性能监控已停止")
    
    def _monitor_loop(self, interval: int) -> None:
        """监控循环"""
        while self.monitoring:
            try:
                stats = self.thread_manager.get_performance_stats()
                
                self.logger.info(
                    f"性能统计 - 总任务: {stats['total_tasks']}, "
                    f"已完成: {stats['completed_tasks']}, "
                    f"失败: {stats['failed_tasks']}, "
                    f"平均处理时间: {stats['average_processing_time']:.2f}s, "
                    f"内存使用: {stats['memory_usage']:.2f}MB, "
                    f"CPU使用: {stats['cpu_usage']:.1f}%"
                )
                
                time.sleep(interval)
                
            except Exception as e:
                self.logger.error(f"性能监控时发生错误: {e}")
                time.sleep(interval)